import asyncio
import websockets
import json
import logging
from datetime import datetime

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class WebSocketClient:
    def __init__(self, token, host='127.0.0.1', port=8765):
        self.token = token
        self.ws = None
        self.connected = False
        self.ws_uri = f"ws://{host}:{port}"
        self.command_handlers = {}
        self.config_handlers = []
        self.task = []

    def on_command(self, command_name):
        """命令处理器装饰器"""
        def decorator(func):
            self.command_handlers[command_name] = func
            return func
        return decorator

    def on_config_update(self, handler):
        """注册配置更新处理器"""
        self.config_handlers.append(handler)
        return handler

    async def connect(self):
        """连接到服务器"""
        try:
            self.ws = await websockets.connect(self.ws_uri)
            # 发送身份验证消息
            await self.ws.send(json.dumps({
                'type': 'auth',
                'token': self.token,
                'timestamp': datetime.now().isoformat()
            }))
            self.connected = True
            logger.info("已连接到服务器")
            return True
        except Exception as e:
            logger.error(f"连接服务器失败: {e}")
            return False

    async def heartbeat(self):
        """发送心跳包"""
        while self.connected:
            try:
                if not self.ws or not self.ws.open:
                    break
                await self.ws.send(json.dumps({
                    'type': 'heartbeat',
                    'timestamp': datetime.now().isoformat()
                }))
                await asyncio.sleep(30)  # 每30秒发送一次心跳
            except Exception as e:
                logger.error(f"发送心跳失败: {e}")
                self.connected = False
                break

    async def handle_command(self, command_data):
        """处理接收到的命令"""
        command = command_data.get('command')
        if command in self.command_handlers:
            try:
                await self.command_handlers[command]()
            except Exception as e:
                logger.error(f"执行命令 {command} 失败: {e}")
        else:
            logger.warning(f"未知命令: {command}")

    async def handle_config_update(self, config_data):
        """处理配置更新"""
        config = config_data.get('config', {})
        logger.info(f"收到配置更新: {config}")
        for handler in self.config_handlers:
            try:
                await handler(config)
            except Exception as e:
                logger.error(f"处理配置更新失败: {e}")

    async def message_handler(self):
        """处理接收到的消息"""
        while self.connected:
            try:
                message = await self.ws.recv()
                data = json.loads(message)
                msg_type = data.get('type')
                
                if msg_type == 'command':
                    await self.handle_command(data)
                elif msg_type == 'config':
                    await self.handle_config_update(data)
                
            except Exception as e:
                logger.error(f"处理消息失败: {e}")
                self.connected = False
                break

    async def start(self):
        """启动客户端"""
        if not await self.connect():
            logger.error("无法连接到服务器，客户端退出")
            return False
        # 启动心跳任务
        heartbeat_task = asyncio.create_task(self.heartbeat())
        # 启动消息处理任务
        handler_task = asyncio.create_task(self.message_handler())
        self.tasks = [heartbeat_task, handler_task]
        # 等待任务完成
        try:
            await asyncio.gather(*self.tasks)
        except Exception as e:
            logger.error(f"客户端运行错误: {e}")
        finally:
            self.connected = False
        
        return True

    async def stop(self):
        """停止客户端"""
        self.connected = False
        if self.ws:
            await self.ws.close()
        print("开始停止")
        if self.tasks:
            for task in self.tasks:
                if not task.done():
                    task.cancel() 
                    try:
                        await task  # 等待任务取消
                    except asyncio.CancelledError:
                        logger.info(f"任务 {task} 已取消")
